import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
import logging
from tqdm import tqdm


class DACDataset(Dataset):
    """针对Criteo Display Advertising Challenge (dac-v1)数据集的自定义数据集类"""

    # self.data 原始数据数据
    # self.num_cols_indices 连续类型特征索引
    # self.cat_cols_indices 离散型特征索引
    # self.target_col_index 目标值索引
    # num_scaler 数据处理 可自定义 MinMaxScaler()用于将数据转换为0~1之间的数据
    # self.cat_encoders 离散数据枚举 le进行列表编码后的结果
    # self.cat_dims 每个离散纬度的枚举个数

    def __init__(self, file_path, cat_encoders=None, num_scaler=None, sample_size=None):
        self.file_path = file_path
        self.sample_size = sample_size

        # 读取数据
        logging.info(f"开始读取数据文件: {file_path}")
        self.data = pd.read_csv(
            file_path,
            sep='\t',
            header=None,
            nrows=sample_size,
            low_memory=False
        )
        logging.info(f"数据读取完成，形状: {self.data.shape}")

        # 验证列数
        expected_cols = 39
        if self.data.shape[1] != expected_cols:
            raise ValueError(f"数据列数不匹配: 预期{expected_cols}列，实际{self.data.shape[1]}列")

        # 填充缺失值
        self.data = self.data.fillna({i: 'NaN' for i in range(self.data.shape[1])})

        # 特征索引配置
        self.num_cols_indices = list(range(0, 13))  # 1-13列：整数特征
        self.cat_cols_indices = list(range(13, expected_cols))  # 14列及以后：类别特征
        self.target_col_index = 0  # 训练集第0列是目标值

        self.num_scaler = num_scaler
        self.data[self.num_cols_indices] = self.data[self.num_cols_indices].replace('NaN', 0).infer_objects(
            copy=False)
        self.data[self.num_cols_indices] = self.data[self.num_cols_indices].astype(float)
        self.data[self.num_cols_indices] = self.num_scaler.transform(self.data[self.num_cols_indices])

        self.cat_encoders = cat_encoders
        # 存储每个类别特征的最频繁值（从训练数据中获取）
        self.most_frequent_values = {}

        for idx in tqdm(self.cat_cols_indices, desc="转换类别特征"):
            self.data[idx] = self.data[idx].astype(str)

            # 获取该特征在训练集中的所有已知类别
            known_classes = set(self.cat_encoders[idx + 1].classes_)

            # 找到测试集中的未知类别
            unknown_mask = ~self.data[idx].isin(known_classes)
            unknown_count = unknown_mask.sum()

            if unknown_count > 0:
                # 如果没有缓存最频繁值，则计算并缓存
                if idx not in self.most_frequent_values:
                    # 从编码器的类中获取训练集中的所有类别
                    # 这里假设编码器的classes_顺序与训练数据中的频率无关
                    # 所以我们需要另一种方式获取最频繁值（如果训练时保存了的话）
                    # 或者这里可以简单使用classes_[0]，或者需要额外逻辑
                    # 注意：这部分需要根据你保存最频繁值的方式调整
                    # 这里假设你在训练时已经计算并存储了最频繁值
                    # 如果没有，你可能需要重新计算或调整策略
                    self.most_frequent_values[idx] = self.cat_encoders[idx + 1].classes_[0]

                # 使用最频繁类别替换未知类别
                most_freq = self.most_frequent_values[idx]
                logging.info(f"特征 {idx} 发现 {unknown_count} 个未知类别，使用最频繁类别 '{most_freq}' 替换")
                self.data.loc[unknown_mask, idx] = most_freq

            # 转换类别特征
            self.data[idx] = self.cat_encoders[idx + 1].transform(self.data[idx])

        # 类别特征维度
        self.cat_dims = [len(le.classes_) for le in self.cat_encoders.values()]
        logging.info(
            f"特征处理完成 - 数值特征: {len(self.num_cols_indices)}个, 类别特征: {len(self.cat_cols_indices)}个")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data.iloc[idx]

        # 提取数值特征
        num_features = torch.tensor([row[i] for i in self.num_cols_indices], dtype=torch.float32)

        # 提取类别特征
        cat_features = torch.tensor([row[i] for i in self.cat_cols_indices], dtype=torch.long)

        return cat_features, num_features
